package org.example.provider;


import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.example.entity.ChatMessageEntity;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.LinkedList;
import java.util.List;


/**
 * 智能助手的记忆存储
 */
@Component
public class AiChatAgentMemoryProvider implements ChatMemoryStore {

    @Autowired
    private MongoTemplate mongoTemplate;

    @Override
    public List<ChatMessage> getMessages(Object messageId) {
        Criteria criteria = Criteria.where("userId").is(messageId);
        Query query = new Query(criteria);
        List<ChatMessageEntity> chatMessageEntities = mongoTemplate.find(query, ChatMessageEntity.class);
        if (chatMessageEntities != null && !chatMessageEntities.isEmpty()) {
            ChatMessageEntity chatMessageEntity = chatMessageEntities.get(0);
            String message = chatMessageEntity.getMessage();
            return ChatMessageDeserializer.messagesFromJson(message);
        }
        return new LinkedList<>();
    }

    @Override
    public void updateMessages(Object messageId, List<ChatMessage> list) {
        Criteria criteria = Criteria.where("userId").is(messageId);
        Query query = new Query(criteria);
        String json = ChatMessageSerializer.messagesToJson(list);
        Update update = new Update().set("message", json);
        mongoTemplate.upsert(query, update, ChatMessageEntity.class);
    }

    @Override
    public void deleteMessages(Object messageId) {
        Criteria  criteria = Criteria.where("userId").is(messageId);
        Query query = new Query(criteria);
        mongoTemplate.remove(query, ChatMessageEntity.class);
    }
}
